{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### unitycatalog-litellm E2E example\n",
    "Import and run this notebook in Databricks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9d4f9e83-c840-4f49-8344-78613916d67c",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "%pip install litellm unitycatalog-ai unitycatalog-client unitycatalog-litellm -q -U\n",
    "dbutils.library.restartPython()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a DatabricksFunctionClient\n",
    "Use serverless client here as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9bdec3f7-edf3-42f5-a742-e8f16c3b53fe",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "127720ad-c167-4e58-b74b-294501b5b0a7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.base import set_uc_function_client\n",
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n",
    "\n",
    "client = DatabricksFunctionClient()\n",
    "\n",
    "# set the default uc function client, it will be used for all toolkits\n",
    "set_uc_function_client(client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a UC function for executing python code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f0b16e6b-3c36-4bbf-bc4b-50fc81fe2bdb",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# replace with your own catalog and schema\n",
    "CATALOG = \"main\"\n",
    "SCHEMA = \"default\"\n",
    "\n",
    "\n",
    "def execute_python_code(code: str) -> str:\n",
    "    \"\"\"\n",
    "    Executes the given python code and returns its stdout.\n",
    "\n",
    "    Args:\n",
    "      code: Python code to execute. Remember to print the final result to stdout.\n",
    "    \"\"\"\n",
    "    import sys\n",
    "    from io import StringIO\n",
    "\n",
    "    stdout = StringIO()\n",
    "    sys.stdout = stdout\n",
    "    exec(code)\n",
    "    return stdout.getvalue()\n",
    "\n",
    "\n",
    "function_info = client.create_python_function(\n",
    "    func=execute_python_code, catalog=CATALOG, schema=SCHEMA, replace=True\n",
    ")\n",
    "python_execution_function_name = function_info.full_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a UC function for translation\n",
    "\n",
    "Use Databricks built-in [ai_translate](https://docs.databricks.com/en/sql/language-manual/functions/ai_translate.html) function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "634b0a21-2763-40d3-be1b-73f8ab2cf155",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local_disk0/.ephemeral_nfs/envs/pythonEnv-ac15189f-1e8b-4bba-b9f2-601de000d4e5/lib/python3.10/site-packages/unitycatalog/ai/core/databricks.py:298: UserWarning: The function translate does not have a description. Using Unity Catalog functions that do not have function descriptions limits the functionality for an LLM to understand when it is appropriate to call your function as a tool and how to properly interface with the function. Update your function's description with a verbose entry in the 'comments' parameter to improve the usage characterstics of this function as a tool.\n",
      "  check_function_info(created_function_info)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FunctionInfo(catalog_name='main', comment=None, created_at=1734104011051, created_by='michael.berk@databricks.com', data_type=<ColumnTypeName.STRING: 'STRING'>, external_language=None, external_name=None, full_data_type='STRING', full_name='main.default.translate', function_id='2c244f6f-61b1-4fc3-9bc6-dbe5e8daea0a', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='content', type_text='string', type_name=<ColumnTypeName.STRING: 'STRING'>, position=0, comment='the text to be translated', parameter_default=None, parameter_mode=None, parameter_type=<FunctionParameterType.PARAM: 'PARAM'>, type_interval_type=None, type_json='{\"name\":\"content\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"the text to be translated\"}}', type_precision=0, type_scale=0), FunctionParameterInfo(name='language', type_text='string', type_name=<ColumnTypeName.STRING: 'STRING'>, position=1, comment='the target language code to translate the content to', parameter_default=None, parameter_mode=None, parameter_type=<FunctionParameterType.PARAM: 'PARAM'>, type_interval_type=None, type_json='{\"name\":\"language\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"the target language code to translate the content to\"}}', type_precision=0, type_scale=0)]), is_deterministic=True, is_null_call=None, metastore_id='b169b504-4c54-49f2-bc3a-adf4b128f36d', name='translate', owner='michael.berk@databricks.com', parameter_style=<FunctionInfoParameterStyle.S: 'S'>, properties='{\"sqlConfig.spark.sql.ansi.enabled\":\"true\",\"referredTempFunctionsNames\":\"[]\",\"sqlConfig.spark.sql.streaming.statefulOperator.stateRebalancing.enabled\":\"false\",\"catalogAndNamespace.part.0\":\"main\",\"sqlConfig.spark.sql.legacy.createHiveTableByDefault\":\"false\",\"sqlConfig.spark.sql.streaming.stopTimeout\":\"15s\",\"referredTempViewNames\":\"[]\",\"catalogAndNamespace.part.1\":\"default\",\"sqlConfig.spark.sql.readSideCharPadding\":\"true\",\"sqlConfig.spark.sql.variable.substitute\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiForecast.enabled\":\"false\",\"sqlConfig.spark.sql.sources.default\":\"delta\",\"sqlConfig.spark.sql.hive.convertCTAS\":\"true\",\"sqlConfig.spark.databricks.sql.functions.vectorSearch.enabled\":\"true\",\"referredTempVariableNames\":\"[]\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOnSocketTimeoutException\":\"true\",\"sqlConfig.spark.sql.sources.commitProtocolClass\":\"com.databricks.sql.transaction.directory.DirectoryAtomicCommitProtocol\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOn400TimeoutError\":\"true\",\"catalogAndNamespace.numParts\":\"2\",\"sqlConfig.spark.sql.stableDerivedColumnAlias.enabled\":\"true\",\"sqlConfig.spark.sql.parquet.compression.codec\":\"snappy\",\"sqlConfig.spark.sql.streaming.stateStore.providerClass\":\"com.databricks.sql.streaming.state.RocksDBStateStoreProvider\"}', return_params=None, routine_body=<FunctionInfoRoutineBody.SQL: 'SQL'>, routine_definition='(SELECT ai_translate(content, language))', routine_dependencies=None, schema_name='default', security_type=<FunctionInfoSecurityType.DEFINER: 'DEFINER'>, specific_name='translate', sql_data_access=<FunctionInfoSqlDataAccess.CONTAINS_SQL: 'CONTAINS_SQL'>, sql_path=None, updated_at=1734104011051, updated_by='michael.berk@databricks.com')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate_function_name = f\"{CATALOG}.{SCHEMA}.translate\"\n",
    "sql_body = f\"\"\"CREATE FUNCTION {translate_function_name}(content STRING COMMENT 'the text to be translated', language STRING COMMENT 'the target language code to translate the content to')\n",
    "RETURNS STRING\n",
    "RETURN SELECT ai_translate(content, language)\n",
    "\"\"\"\n",
    "client.create_function(sql_function_body=sql_body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1fe93745-28a2-4f2c-980e-684418d0f392",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local_disk0/.ephemeral_nfs/envs/pythonEnv-ac15189f-1e8b-4bba-b9f2-601de000d4e5/lib/python3.10/site-packages/unitycatalog/ai/core/databricks.py:581: UserWarning: The function translate does not have a description. Using Unity Catalog functions that do not have function descriptions limits the functionality for an LLM to understand when it is appropriate to call your function as a tool and how to properly interface with the function. Update your function's description with a verbose entry in the 'comments' parameter to improve the usage characterstics of this function as a tool.\n",
      "  check_function_info(function_info)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value='¿Qué es Databricks?', truncated=None)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    translate_function_name, {\"content\": \"What is Databricks?\", \"language\": \"es\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2675f057-e152-45fe-8a38-b710cf23b3d4",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/spark-ac15189f-1e8b-4bba-b9f2-60/.ipykernel/14958/command-1819603593741393-926642350:3: UserWarning: A custom validator is returning a value other than `self`.\n",
      "Returning anything other than `self` from a top level model validator isn't supported when validating via `__init__`.\n",
      "See the `model_validator` docs (https://docs.pydantic.dev/latest/concepts/validators/#model-validators) for more details.\n",
      "  toolkit = UCFunctionToolkit(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[LiteLLMTool(name='main__default__execute_python_code', description='Executes the given python code and returns its stdout.', tool={'type': 'function', 'function': {'name': 'main__default__execute_python_code', 'strict': True, 'parameters': {'properties': {'code': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'description': 'Python code to execute. Remember to print the final result to stdout.', 'title': 'Code'}}, 'title': 'main__default__execute_python_code__params', 'type': 'object', 'additionalProperties': False, 'required': ['code']}, 'description': 'Executes the given python code and returns its stdout.'}}),\n",
       " LiteLLMTool(name='main__default__translate', description='', tool={'type': 'function', 'function': {'name': 'main__default__translate', 'strict': True, 'parameters': {'properties': {'content': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'description': 'the text to be translated', 'title': 'Content'}, 'language': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'description': 'the target language code to translate the content to', 'title': 'Language'}}, 'title': 'main__default__translate__params', 'type': 'object', 'additionalProperties': False, 'required': ['content', 'language']}, 'description': ''}})]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from unitycatalog.ai.litellm.toolkit import UCFunctionToolkit\n",
    "\n",
    "toolkit = UCFunctionToolkit(\n",
    "    function_names=[python_execution_function_name, translate_function_name]\n",
    ")\n",
    "tools = toolkit.tools\n",
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use the tools in LiteLLM Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "102272ba-fd86-414a-9bed-d08502c2ff50",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ModelResponse(id='chatcmpl-Ae2IAKtJJzP0P3ynKkCudoBhNcJpz', created=1734105294, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_6fc10e10eb', choices=[Choices(finish_reason='tool_calls', index=0, message=Message(content=None, role='assistant', tool_calls=[ChatCompletionMessageToolCall(function=Function(arguments='{\"content\":\"cerdo hormiguero\",\"language\":\"en\"}', name='main__default__translate'), id='call_Vz5ZG9JUgudEv5e1gM4o2wVg', type='function')], function_call=None))], usage=Usage(completion_tokens=25, prompt_tokens=152, total_tokens=177, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier=None)\n"
     ]
    }
   ],
   "source": [
    "import litellm\n",
    "\n",
    "# Define your request\n",
    "question = \"What is 'cerdo hormiguero' in english?\"\n",
    "messages = [{\"role\": \"user\", \"content\": question}]\n",
    "\n",
    "# Show the response\n",
    "response = litellm.completion(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=messages,\n",
    "    tools=tools,\n",
    "    tool_choice=\"auto\",  # auto is default, but we'll be explicit\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Leverage `generate_tool_call_messages` to run the UC Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "569a79fa-5ad5-44b0-a98e-0837c45c7e12",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local_disk0/.ephemeral_nfs/envs/pythonEnv-ac15189f-1e8b-4bba-b9f2-601de000d4e5/lib/python3.10/site-packages/unitycatalog/ai/core/databricks.py:581: UserWarning: The function translate does not have a description. Using Unity Catalog functions that do not have function descriptions limits the functionality for an LLM to understand when it is appropriate to call your function as a tool and how to properly interface with the function. Update your function's description with a verbose entry in the 'comments' parameter to improve the usage characterstics of this function as a tool.\n",
      "  check_function_info(function_info)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ModelResponse(id='chatcmpl-Ae2IDXCqk5xJPOJzhGjiY7MN6ykGN', created=1734105297, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_6fc10e10eb', choices=[Choices(finish_reason='stop', index=0, message=Message(content='The term \"cerdo hormiguero\" translates to \"ant bear\" in English, which is commonly known as the \"anteater.\"', role='assistant', tool_calls=None, function_call=None))], usage=Usage(completion_tokens=28, prompt_tokens=57, total_tokens=85, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier=None)\n"
     ]
    }
   ],
   "source": [
    "from unitycatalog.ai.litellm.utils import generate_tool_call_messages\n",
    "\n",
    "tool_messages = generate_tool_call_messages(\n",
    "    response=response, client=client, conversation_history=messages\n",
    ")\n",
    "\n",
    "response_2 = litellm.completion(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=tool_messages,\n",
    ")\n",
    "response_2"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "computePreferences": null,
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 2
   },
   "notebookName": "UC LiteLLM",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
